package cn.codecrab.ai.langchain4j.controller;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.ClassPathDocumentLoader;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import jakarta.annotation.Resource;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.nio.file.FileSystems;
import java.util.List;

/**
 * @author 王刚
 * @since 2025年05月11日 17:57
 */
@RestController
@RequestMapping("/embedding")
public class EmbeddingController {

    @Resource
    private EmbeddingModel embeddingModel;

    @Resource
    private EmbeddingStore<TextSegment> esEmbeddingStore;

    @Resource
    private EmbeddingStore<TextSegment> pineconeEmbeddingStore;

    @GetMapping
    public String embedding() {
        Response<Embedding> response = embeddingModel.embed("你好呀");
        return response.content().vector().length + "\n" + response;
    }

    @GetMapping("/init")
    public void initEmbedding() {
        List<Document> documents = ClassPathDocumentLoader.loadDocuments("knowledge/hospital", FileSystems.getDefault().getPathMatcher("glob:*.md"));
        EmbeddingStoreIngestor.builder().embeddingModel(embeddingModel).embeddingStore(pineconeEmbeddingStore).build().ingest(documents);
    }

    @GetMapping("/store")
    public void embeddingStore() {
        TextSegment segment = TextSegment.from("我喜欢打羽毛球");
        esEmbeddingStore.add(embeddingModel.embed(segment).content(), segment);
        pineconeEmbeddingStore.add(embeddingModel.embed(segment).content(), segment);

        segment = TextSegment.from("我喜欢打篮球");
        esEmbeddingStore.add(embeddingModel.embed(segment).content(), segment);
        pineconeEmbeddingStore.add(embeddingModel.embed(segment).content(), segment);

        segment = TextSegment.from("今天天气很好");
        esEmbeddingStore.add(embeddingModel.embed(segment).content(), segment);
        pineconeEmbeddingStore.add(embeddingModel.embed(segment).content(), segment);
    }

    @GetMapping("/search")
    public void embeddingSearch() {
        TextSegment segment = TextSegment.from("你喜欢什么运动");

        EmbeddingSearchRequest searchRequest = EmbeddingSearchRequest.builder()
            .queryEmbedding(embeddingModel.embed(segment).content())
            .maxResults(1)
            .minScore(0.8)
            .build();

        // System.out.println(esEmbeddingStore.search(searchRequest)); // Missing required property 'KnnQuery.k'
        System.out.println("===============================");
        System.out.println(pineconeEmbeddingStore.search(searchRequest).matches());

        System.out.println("\n\n\n");

        segment = TextSegment.from("今天天气怎么样");

        searchRequest = EmbeddingSearchRequest.builder()
            .queryEmbedding(embeddingModel.embed(segment).content())
            .maxResults(1)
            .minScore(0.8)
            .build();

        // System.out.println(esEmbeddingStore.search(searchRequest)); // Missing required property 'KnnQuery.k'
        System.out.println("===============================");
        System.out.println(pineconeEmbeddingStore.search(searchRequest).matches());
    }

}
